package com.xiyuan.smartutils.crypto;

import com.xiyuan.smartutils.Asserts;
import com.xiyuan.smartutils.Hexs;
import com.xiyuan.smartutils.Systems;
import com.xiyuan.smartutils.constants.CodeConstants;
import com.xiyuan.smartutils.extend.KV;
import com.xiyuan.smartutils.reflect.Unsafe;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.nio.charset.Charset;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

/**
 * RSA非对称加解密
 *
 * @version v1.0.0 @author lgz 2018-8-22 新建与整理
 */
public class RSA implements CodeConstants {
    static {
        Unsafe.init();
    }
    /**********************************************************************************/
    //生成密钥对
    /**********************************************************************************/
    
    /**
     * 生成一对公钥和私钥，第一个为公钥，第二个为私钥，默认1024
     */
    public static KV<String, String> buildKeyPair() {
        return buildKeyPairBase64(KiB);
    }
    
    /**
     * 生成一对公钥和私钥
     *
     * @param keySize 密钥长度，如512/1024/2048
     * @return 第一个为公钥，第二个为私钥
     */
    public static KV<String, String> buildKeyPairBase64(int keySize) {
        try {
            KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(_RSA_);
            keyPairGen.initialize(keySize);
            KeyPair keyPair = keyPairGen.generateKeyPair();
            
            String publicKey = Base64.encode(keyPair.getPublic().getEncoded());
            String privateKey = Base64.encode(keyPair.getPrivate().getEncoded());
            return new KV<String, String>(publicKey, privateKey);
        }
        catch (NoSuchAlgorithmException e) {
            throw Asserts.exception(e);
        }
    }
    
    /**
     * 生成一对公钥和私钥
     *
     * @param keySize 密钥长度，如512/1024/2048
     * @return 第一个为公钥，第二个为私钥
     */
    public static KV<byte[], byte[]> buildKeyPair(int keySize) {
        try {
            KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(_RSA_);
            keyPairGen.initialize(keySize);
            KeyPair keyPair = keyPairGen.generateKeyPair();
            
            return new KV<>(keyPair.getPublic().getEncoded(), keyPair.getPrivate().getEncoded());
        }
        catch (NoSuchAlgorithmException e) {
            throw Asserts.exception(e);
        }
    }
    
    /**********************************************************************************/
    //标准加解密
    /**********************************************************************************/
    
    /**
     * 加密（使用公钥加密），支持常用的512/1024/2048/3072/4096长度适配
     *
     * @param src       原数据
     * @param publicKey 公钥字符串（BASE64）
     * @return 加密后数据
     * @throws Exception
     */
    public static byte[] encrypt(byte[] src, String publicKey) throws Exception {
        return encrypt(src, Base64.decode(publicKey));
    }
    
    /**
     * 加密（使用公钥加密），支持常用的512/1024/2048/3072/4096长度适配
     *
     * @param src       原数据
     * @param publicKey 公钥字节数组
     * @return 加密后数据
     * @throws Exception 异常信息
     */
    public static byte[] encrypt(byte[] src, byte[] publicKey) throws Exception {
        int keySize = 0;
        switch (publicKey.length) {
            case 94:
                keySize = 512;
                break;//343
            case 162:
                keySize = 1024;
                break;//634
            case 294:
                keySize = 2048;
                break;//1217
            case 422:
                keySize = 3072;
                break;//1793
            case 550:
                keySize = 4096;
                break;//2376
        }
        
        Asserts.as(keySize != 0
                   ? null
                   : "无法匹配密钥长度，请调用encrypt(byte[] src, byte[] publicKey, int keySize)方法");
        return encrypt(src, publicKey, keySize);
    }
    
    /**
     * 加密（使用公钥加密）
     *
     * @param src       原数据
     * @param publicKey 公钥字节数组
     * @param keySize   密钥长度，如512/1024/2048
     * @return 加密后数据
     * @throws Exception 可能的异常
     */
    public static byte[] encrypt(byte[] src, byte[] publicKey, int keySize) throws Exception {
        //1.组装公钥
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(publicKey);
        KeyFactory keyFactory = KeyFactory.getInstance(_RSA_);
        PublicKey key = keyFactory.generatePublic(keySpec);
        
        //2.初始化密码器
        Cipher cipher = Cipher.getInstance(_RSA_);
        cipher.init(Cipher.ENCRYPT_MODE, key);
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        
        //3.加密时分块加密，每块支持大小为keySize / 8 - 11
        int length = src.length;
        int chunkSize = keySize / 8 - 11;
        
        for (int i = 0; i * chunkSize < length; i++) {
            int count = length > (i + 1) * chunkSize ? chunkSize : length - i * chunkSize;
            out.write(cipher.doFinal(src, i * chunkSize, count));
        }
        
        return out.toByteArray();
    }
    
    /**
     * 解密（使用私钥加密），支持常用的512/1024/2048/3072/4096长度适配
     *
     * @param src        原数据
     * @param privateKey 私钥字符串（BASE64）
     * @return 加密后数据
     * @throws Exception 可能的异常
     */
    public static byte[] decrypt(byte[] src, String privateKey) throws Exception {
        return decrypt(src, Base64.decode(privateKey));
    }
    
    /**
     * 解密（使用私钥加密），支持常用的512/1024/2048/3072/4096长度适配
     *
     * @param src        原数据
     * @param privateKey 私钥字节数组
     * @return 加密后数据
     * @throws Exception 可能的异常
     */
    public static byte[] decrypt(byte[] src, byte[] privateKey) throws Exception {
        int keyLen = privateKey.length;
        int keySize = 0;
    
        if (keyLen >= 330 && keyLen <= 360)//342-347
        {keySize = 512;}
        else if (keyLen >= 620 && keyLen <= 650)//631-637
        {keySize = 1024;}
        else if (keyLen >= 1200 && keyLen <= 1230)//1214-1220
        {keySize = 2048;}
        else if (keyLen >= 1780 && keyLen <= 1710)//1792-1796
        {keySize = 3072;}
        else if (keyLen >= 2360 && keyLen <= 2390)//2373-2376
        {keySize = 4096;}
        
        Asserts.as(keySize != 0
                   ? null
                   : "无法匹配密钥长度，请调用decrypt(byte[] src, byte[] keyBytes, int keySize)方法");
        return decrypt(src, privateKey, keySize);
    }
    
    /**
     * 解密（使用私钥加密）
     *
     * @param src        原数据
     * @param privateKey 私钥字节数组
     * @param keySize    密钥长度，如512/1024/2048
     * @return 加密后数据
     * @throws Exception 可能的异常
     */
    public static byte[] decrypt(byte[] src, byte[] privateKey, int keySize) throws Exception {
        //1.组装私钥
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKey);
        KeyFactory keyFactory = KeyFactory.getInstance(_RSA_);
        PrivateKey key = keyFactory.generatePrivate(keySpec);
        
        
        //2.初始化密码器
        Cipher cipher = Cipher.getInstance(_RSA_);
        cipher.init(Cipher.DECRYPT_MODE, key);
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        
        //3.解密时分块加密，每块支持大小为keySize / 8
        int length = src.length;
        int chunkSize = keySize / 8;
        
        for (int i = 0; i * chunkSize < length; i++) {
            int count = length > (i + 1) * chunkSize ? chunkSize : length - i * chunkSize;
            out.write(cipher.doFinal(src, i * chunkSize, count));
        }
        
        return out.toByteArray();
    }
    
    /**********************************************************************************/
    //字符串加解密
    /**********************************************************************************/
    
    /**
     * 字符串加密
     */
    public static byte[] encrypt(String src, Charset charset, String publicKey) throws Exception {
        return encrypt(src.getBytes(charset), publicKey);
    }
    
    /**
     * 字符串加密
     */
    public static byte[] encrypt(String src, String encoding, String publicKey) throws Exception {
        return encrypt(src, Charset.forName(encoding), publicKey);
    }
    
    /**
     * 字符串解密
     */
    public static String decrypt(byte[] src, String privateKey, Charset charset) throws Exception {
        return new String(decrypt(src, privateKey), charset);
    }
    
    /**
     * 字符串解密
     */
    public static String decrypt(byte[] src, String privateKey, String encoding) throws Exception {
        return decrypt(src, privateKey, Charset.forName(encoding));
    }
    
    /**********************************************************************************/
    //字符串BASE64加解密
    /**********************************************************************************/
    
    /**
     * 字符串加密
     * @param src        需要加密的字符串
     * @param publicKey  RSA公钥 base64 字符串
     * @param encoding   字符串编码
     * @return      返回加密后的base64 字符串
     * @throws Exception 异常信息
     */
    public static String encode(String src, String encoding, String publicKey) throws Exception {
        return Base64.encode(encrypt(src.getBytes(encoding), publicKey));
    }
    
    /**
     * 字符串加密
     * @param src        需要加密的字符串
     * @param publicKey  RSA公钥 base64 字符串
     * @return      返回加密后的base64 字符串
     * @throws Exception 异常信息
     */
    public static String encodeUTF8(String src, String publicKey) throws Exception {
        return Base64.encode(encrypt(src.getBytes(_UTF_8_C_), publicKey));
    }
    
    /**
     * 字符串解密
     * @param src        RSA加密后的base64 字符串
     * @param privateKey RSA 私钥 base64 字符串
     * @param encoding   返回的字符串编码
     * @return 返回解密后以指定编码的字符串
     * @throws Exception 异常信息
     */
    public static String decode(String src, String privateKey, String encoding) throws Exception {
        return new String(decrypt( Base64.decode(src), privateKey), encoding);
    }
    
    /**
     * 字符串解密
     * @param src        RSA加密后的base64 字符串
     * @param privateKey RSA 私钥 base64 字符串
     * @return 返回解密后以UTF8 编码的字符串
     * @throws Exception 异常信息
     */
    public static String decodeUTF8(String src, String privateKey) throws Exception {
        return new String(decrypt( Base64.decode(src), privateKey), _UTF_8_C_);
    }
    
    
    /**********************************************************************************/
    //字符串HEX加解密
    /**********************************************************************************/
    
    /**
     * 字符串加密
     */
    public static String encryptHexUTF8(String src, String publicKey) throws Exception {
        return encryptHex(src, _UTF_8_C_, publicKey);
    }
    
    /**
     * 字符串加密
     */
    public static String encryptHex64(String src, String encoding, String publicKey) throws Exception {
        return encryptHex(src, Charset.forName(encoding), publicKey);
    }
    
    /**
     * 字符串加密
     */
    public static String encryptHex(String src, Charset charset, String publicKey) throws Exception {
        return Hexs.toHexString(encrypt(src.getBytes(charset), publicKey));
    }
    
    /**
     * 字符串解密
     */
    public static String decryptHexUTF8(String src, String privateKey) throws Exception {
        return decryptHex(src, privateKey, _UTF_8_C_);
    }
    
    /**
     * 字符串解密
     */
    public static String decryptHex(String src, String privateKey, String encoding) throws Exception {
        return decryptHex(src, privateKey, Charset.forName(encoding));
    }
    
    /**
     * 字符串解密
     */
    public static String decryptHex(String src, String privateKey, Charset charset) throws Exception {
        byte[] bytes = Hexs.toBytes(src);
        return new String(decrypt(bytes, privateKey), charset);
    }
}
